# % 该函数主要用于构成trunk的长度和索引
# % 在偏振模色散中，由于其会发生随机的相位旋转，将传输长度分为多段段距离的保偏光纤
# % 在这一小段中传输，传输的光脉冲相位和大小不受PMD的影响，但不同之间会发生随机相位旋转和差分群时延
#
# %函数输出：
# %L_trunk：
# %       在step_size内传输时，所经过的trunk的长度，如果在同一个trunk中传输，则L_trunk = step_size
# %       如果一个 step_size被分成在两个trunk中传输，则 L_trunk将表示一个向量，对应的包含两个trunk中分别的传输长度。
# %Index_trunk：
# %       对应的trunk索引值，若在一个trunk中传输，则Index_trunk表示当前传输的trunk段，例如第2个trunk中
# %       若在两个trunk中传输，如经过了2，3两个，则Index_trunk = [2,3]
#
# %函数输入：
# %Final_distance：
# %       分步传输后每一次的传输终点的值，单位为km
# % step_size：
# %       SSFM中每一步的步长值
# % trunk_num：
# %       所设置的trunk数，即将一整个长光纤长度分为多少个短的保偏光纤
#
from numpy import ceil, floor, diff
import numpy as np

def Trunk_Composition(Final_distance, step_size, Fiber_para, PMD_para):

    trunk_length = PMD_para["L_corr"]

# ??????????????????????????????????????????????????????????????????????
    TOL0 = 1e-12        # [km] tolerance to claim a step of zero length

    z_start = Final_distance - step_size     # 每个step_size对应的起点 [km]
    z_end = Final_distance     # 每个step_size结束对应的终点 [km]
    if z_end > Fiber_para["L"]:
        z_end = Fiber_para["L"]

# 设置第一个trunk的索引值为1，因为floor(10)+1=11, but ceil(10)=10，故需保证第一个索引值
    index_start = floor(z_start/trunk_length)+1     # 相应trunk的起始索引，即从第几个trunk开始

    index_end = ceil(z_end/trunk_length)      # 相应trunk的结束索引

# 若该step_size在同一段trunk中传输，则index_start = index_end，若不同，即表示该step_size跨越了两个trunk传输
    index = np.arange(index_start, index_end+1)         # 所经历的trunk索引
    # print(index)

    if index_start == index_end:        # 在同一段trunk内进行传输
        z_split = step_size
    else:        # 在不同的trunk中传输
        zmid = trunk_length*index[0:-1]     # 两段trunk的交点
        z_split = [zmid[0]-z_start, diff(zmid), z_end-zmid[-1]]
        # print("z_split:", z_split)
        z_split = [item if not isinstance(item, np.ndarray) else item.item() for item in z_split if item != 0]

    z_split = np.array(z_split)
    # print("z_split shape:", z_split.shape)
    # z_split = z_split
    # 由于step-size被分为了两个trunk进行传输，故dzsplit[，]分别表示在前一个trunk的传输距离和在后一个trunk内传输的距离

    mask = np.abs(z_split) > TOL0
    PMD_para['L_trunk'] = z_split[mask]          # remove zero-length steps
    # print("Mask shape:", mask.shape)
    # print("Index shape:", index.shape)
    # assert mask.shape == index.shape
    PMD_para['Index_trunk'] = index[mask]         # remove zero-length steps
    # print('----------------PMD参数构造------------------')
    # print('\n')
    # print('总传输距离 = %.4f\n' % Final_distance)
    # print('开始长度 = %.4f 结束长度 = %.4f 构造的Trunk长度 = %.4f\n' %
    #       (z_start, z_end, trunk_length))
    # print('开始索引 = %d 结束索引 = %d\n' % (index_start, index_end))
    # # print('L_trunk = ')
    # # print('%.4f ' % [PMD_para['L_trunk'][i] for i in range(2) ])
    # print(f'L_trunk = {PMD_para["L_trunk"]}')
    # print('\n')
    # # print('z_split = ')
    # # print('%.4f ',  z_split)
    # print(f'z_split = {z_split}')
    # print('\n')
    # print('---------------------------------------------------')
    # print('\n')
    return PMD_para


if __name__ == '__main__':
    Final_distance = 1
    step_size = 1
    Fiber_para = {"L": 80}
    PMD_para = {"L_corr": 0.8}
    PMD_para = Trunk_Composition(Final_distance=Final_distance,
                                 step_size=step_size, Fiber_para=Fiber_para, PMD_para=PMD_para)
    print(PMD_para)
